package com.br.aiops.storage;

import com.alibaba.fastjson.JSONObject;
import com.br.aiops.counter.AlertLogCounter;
import com.br.aiops.counter.Counter;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
import org.springframework.stereotype.Component;

import java.util.Collection;
import java.util.List;
import java.util.stream.Collectors;

@Component
public class KafkaRebalanceListener implements ConsumerAwareRebalanceListener {

    private static final Logger LOGGER = LoggerFactory.getLogger(KafkaRebalanceListener.class);

    @Autowired
    private Counter counter;

    @Autowired
    private PersistentCounterService persistentCounterService;

    @Value("${spring.kafka.consumer.topic}")
    private String topic;

    @Override
    public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
        if (partitions.stream().anyMatch(partition -> partition.topic().equals(topic))) {
            LOGGER.info("移走的partitions：{}", partitions.stream().map(TopicPartition::toString).collect(Collectors.joining(",")));
            LOGGER.info("持久化当前的counter：{}", JSONObject.toJSONString(counter.getCounter()));
            //先持久化目前所有数据
            persistentCounterService.storage();
            List<Integer> partitionList = partitions.stream().filter(partition -> partition.topic().equals(topic))
                    .map(TopicPartition::partition).collect(Collectors.toList());
            //清除counter统计，仅清除log-alert的
            List<String> removeKeys = counter.getCounter().values().stream()
                    .filter(alertLogCounter -> partitionList.contains(alertLogCounter.getAlertLog().getPartition()))
                    .map(AlertLogCounter::counterKey)
                    .collect(Collectors.toList());
            removeKeys.forEach(key -> counter.getCounter().remove(key));
        }
    }

    @Override
    public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
        if (partitions.stream().anyMatch(partition -> partition.topic().equals(topic))) {
            LOGGER.info("新指派的partitions：{}", partitions.stream()
                    .map(TopicPartition::toString).collect(Collectors.joining(",")));
            //新指派来的partitions，需要做load
            try {
                //仅加载log-alert的
                partitions.stream().filter(partition -> partition.topic().equals(topic))
                        .forEach(partition -> {
                            try {
                                persistentCounterService.load(partition.partition());
                            } catch (Exception e) {
                                LOGGER.error("加载partition:{}持久化数据失败！", partition, e);
                            }
                        });
                counter.ready();
                LOGGER.info("新的counter为：{}", counter.getCounter());
            } catch (Exception e) {
                LOGGER.error("加载持久化数据失败！", e);
            }
        }
    }
}
